#!/usr/bin/env python3

"""
Neural-SLAM 完整训练和评估工作流程
"""

import os
import sys
import time
import subprocess
import argparse
from datetime import datetime

def run_command(cmd, description="", log_file=None):
    """运行命令并记录输出"""
    print(f"\n{'='*60}")
    print(f"🚀 {description}")
    print(f"📝 命令: {cmd}")
    print(f"{'='*60}")
    
    start_time = time.time()
    
    try:
        if log_file:
            with open(log_file, 'a') as f:
                f.write(f"\n{'='*60}\n")
                f.write(f"时间: {datetime.now()}\n")
                f.write(f"描述: {description}\n")
                f.write(f"命令: {cmd}\n")
                f.write(f"{'='*60}\n")
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"✅ 成功完成: {description}")
            if result.stdout:
                print(f"输出:\n{result.stdout[-500:]}")  # 只显示最后500字符
        else:
            print(f"❌ 失败: {description}")
            print(f"错误: {result.stderr}")
            return False
            
        if log_file:
            with open(log_file, 'a') as f:
                f.write(f"返回码: {result.returncode}\n")
                f.write(f"输出: {result.stdout}\n")
                if result.stderr:
                    f.write(f"错误: {result.stderr}\n")
                f.write(f"\n")
        
        elapsed = time.time() - start_time
        print(f"⏱️  耗时: {elapsed:.2f}秒")
        return True
        
    except Exception as e:
        print(f"❌ 执行命令时出错: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Neural-SLAM 完整工作流程')
    parser.add_argument('--baseline-episodes', type=int, default=50, 
                       help='基线模型训练episode数')
    parser.add_argument('--improved-episodes', type=int, default=100, 
                       help='改进模型训练episode数')
    parser.add_argument('--skip-baseline', action='store_true',
                       help='跳过基线训练（如果已有baseline数据）')
    parser.add_argument('--skip-improved', action='store_true',
                       help='跳过改进模型训练')
    parser.add_argument('--only-eval', action='store_true',
                       help='仅运行评估和指标提取')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs("./tmp/", exist_ok=True)
    os.makedirs("./results/", exist_ok=True)
    log_file = f"./results/workflow_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
    
    print("🧪 Neural-SLAM 完整训练和评估工作流程")
    print("=" * 60)
    print(f"📅 开始时间: {datetime.now()}")
    print(f"📋 配置:")
    print(f"   基线训练episodes: {args.baseline_episodes}")
    print(f"   改进模型训练episodes: {args.improved_episodes}")
    print(f"   跳过基线训练: {args.skip_baseline}")
    print(f"   跳过改进模型训练: {args.skip_improved}")
    print(f"   仅评估模式: {args.only_eval}")
    print(f"📝 日志文件: {log_file}")
    
    success_count = 0
    total_steps = 0
    
    if not args.only_eval:
        # 步骤1: 基线模型训练
        if not args.skip_baseline:
            total_steps += 1
            cmd = f"python main.py --exp_name baseline_original --num_episodes {args.baseline_episodes} --train_slam 1 --train_global 1 --train_local 1"
            if run_command(cmd, "步骤1: 训练原始Neural-SLAM基线模型", log_file):
                success_count += 1
        
        # 步骤2: 基线性能评估
        total_steps += 1
        cmd = f"python main.py --exp_name baseline_eval --num_episodes 20 --eval 1 --load_slam ./tmp/models/baseline_original/ --load_global ./tmp/models/baseline_original/ --load_local ./tmp/models/baseline_original/"
        if run_command(cmd, "步骤2: 评估基线模型性能", log_file):
            success_count += 1
        
        # 步骤3: 提取基线指标
        total_steps += 1
        cmd = "python extract_baseline_metrics.py --log_dir ./tmp/ --experiment baseline"
        if run_command(cmd, "步骤3: 提取基线性能指标", log_file):
            success_count += 1
        
        # 步骤4: 训练改进模型
        if not args.skip_improved:
            total_steps += 1
            cmd = f"python main.py --exp_name improved_model --num_episodes {args.improved_episodes} --train_slam 1 --train_global 1 --train_local 1"
            if run_command(cmd, "步骤4: 训练改进的Neural-SLAM模型", log_file):
                success_count += 1
        
        # 步骤5: 改进模型评估
        total_steps += 1
        cmd = f"python main.py --exp_name improved_eval --num_episodes 20 --eval 1 --load_slam ./tmp/models/improved_model/ --load_global ./tmp/models/improved_model/ --load_local ./tmp/models/improved_model/"
        if run_command(cmd, "步骤5: 评估改进模型性能", log_file):
            success_count += 1
        
        # 步骤6: 提取改进模型指标
        total_steps += 1
        cmd = "python extract_baseline_metrics.py --log_dir ./tmp/ --experiment improved"
        if run_command(cmd, "步骤6: 提取改进模型性能指标", log_file):
            success_count += 1
    
    # 步骤7: 生成对比报告
    total_steps += 1
    cmd = "python generate_paper_results.py --baseline_dir ./tmp/ --improved_dir ./tmp/ --output_dir ./results/"
    if run_command(cmd, "步骤7: 生成性能对比报告和图表", log_file):
        success_count += 1
    
    # 步骤8: 创建最终结果总结
    total_steps += 1
    if create_summary_report(log_file):
        success_count += 1
        print("✅ 步骤8: 创建最终结果总结")
    else:
        print("❌ 步骤8: 创建最终结果总结失败")
    
    # 输出最终结果
    print("\n" + "=" * 60)
    print("📊 工作流程完成统计")
    print("=" * 60)
    print(f"✅ 成功步骤: {success_count}/{total_steps}")
    print(f"📅 结束时间: {datetime.now()}")
    
    if success_count == total_steps:
        print("🎉 所有步骤成功完成！")
        print("📁 结果文件位置:")
        print("   - 训练日志: ./tmp/")
        print("   - 性能指标: ./results/")
        print("   - 对比图表: ./results/")
        print("   - 工作流程日志: " + log_file)
    else:
        print("⚠️  部分步骤失败，请检查日志文件")
        return 1
    
    return 0

def create_summary_report(log_file):
    """创建最终结果总结"""
    try:
        summary_file = f"./results/summary_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
        
        with open(summary_file, 'w', encoding='utf-8') as f:
            f.write("# Neural-SLAM 实验结果总结\n\n")
            f.write(f"**生成时间**: {datetime.now()}\n\n")
            
            f.write("## 实验概述\n\n")
            f.write("本次实验对比了原始Neural-SLAM模型和改进版本的性能。\n\n")
            
            f.write("## 改进点\n\n")
            f.write("1. **注意力机制**: 引入Multi-Head Self-Attention替代LSTM\n")
            f.write("2. **现代视觉编码器**: 使用EfficientNet作为骨干网络\n")
            f.write("3. **混合精度训练**: 利用FP16加速训练\n")
            f.write("4. **AI2-THOR环境**: 更真实的物理仿真环境\n\n")
            
            f.write("## 文件结构\n\n")
            f.write("```\n")
            f.write("./tmp/\n")
            f.write("├── models/\n")
            f.write("│   ├── baseline_original/     # 基线模型权重\n")
            f.write("│   ├── baseline_eval/         # 基线评估结果\n")
            f.write("│   ├── improved_model/        # 改进模型权重\n")
            f.write("│   └── improved_eval/         # 改进模型评估结果\n")
            f.write("└── logs/                      # 训练日志\n\n")
            f.write("./results/\n")
            f.write("├── performance_comparison.png  # 性能对比图\n")
            f.write("├── metrics_table.tex          # LaTeX表格\n")
            f.write("├── training_curves.png        # 训练曲线\n")
            f.write("└── summary_report_*.md        # 本文件\n")
            f.write("```\n\n")
            
            f.write("## 关键指标\n\n")
            f.write("- **成功率 (Success Rate)**: 到达目标的episode比例\n")
            f.write("- **SPL (Success weighted by Path Length)**: 路径长度加权的成功率\n")
            f.write("- **探索效率**: 单位时间内探索的面积\n")
            f.write("- **地图质量**: 构建地图的准确性\n\n")
            
            f.write("## 使用说明\n\n")
            f.write("1. 查看 `./results/performance_comparison.png` 了解性能对比\n")
            f.write("2. 使用 `./results/metrics_table.tex` 在论文中插入结果表格\n")
            f.write("3. 参考训练日志进行进一步分析\n\n")
            
            f.write("## 下一步工作\n\n")
            f.write("- 在更多场景中验证模型性能\n")
            f.write("- 尝试其他先进的视觉编码器\n")
            f.write("- 优化超参数配置\n")
            f.write("- 扩展到真实机器人平台\n")
        
        print(f"📄 创建总结报告: {summary_file}")
        return True
        
    except Exception as e:
        print(f"❌ 创建总结报告失败: {e}")
        return False

if __name__ == "__main__":
    sys.exit(main())
